#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-05-13

import platform
import subprocess
import sys
import os
import argparse
import shutil
import time
from typing import Dict, Optional
from pathlib import Path


node_dict = {
    'pcam': {
        'info': 'Python Camera Reading Node',
        'app': 'spirecv/dataloader/CameraNode.py',
        'plat': []
    },
    'ccam': {
        'info': 'C++ Camera Reading Node',
        'app': 'spirecv/dataloader/CameraCppNode/build/SV2CameraNode',
        'plat': []
    },
    'pcamcalib': {
        'info': 'Monocular Camera Calibration Node',
        'app': 'spirecv/calib/CameraCalibrationNode.py',
        'plat': []
    },
    'pcocoloader': {
        'info': 'COCO Format Dataset Loading Node',
        'app': 'spirecv/dataloader/COCODatasetLoaderNode.py',
        'plat': []
    },
    'pimreader': {
        'info': 'Image Folder Reading Node',
        'app': 'spirecv/dataloader/ImgDirReaderNode.py',
        'plat': []
    },
    'pbinreader': {
        'info': 'Points Bin Folder Reading Node',
        'app': 'spirecv/dataloader/LidarBinDirReaderNode.py',
        'plat': []
    },
    'pspirereader': {
        'info': 'Spire Format Dataset Loading Node',
        'app': 'spirecv/dataloader/SpireFormatReaderNode.py',
        'plat': []
    },
    'pnusloader': {
        'info': 'NuScenes Dataset Loading Node',
        'app': 'spirecv/dataloader/NuScenesDatasetLoaderNode.py',
        'plat': []
    },
    'pstmreader': {
        'info': 'Python Stream Reading Node',
        'app': 'spirecv/dataloader/StreamReaderNode.py',
        'plat': []
    },
    'pvid': {
        'info': 'Video File Reading Node',
        'app': 'spirecv/dataloader/VideoFileNode.py',
        'plat': []
    },
    'p3lidar2bev': {
        'info': 'Concat 3 Lidar and Transfor to BirdView Node',
        'app': 'spirecv/dataloader/ThreeLidar2BirdviewNode.py',
        'plat': ['cuda']
    },
    'p3lidar2sms': {
        'info': 'Converts 3 ROS2 lidar messages into SMS messages Node',
        'app': 'spirecv/fromros2/ROS2ThreeLidar2SMSNode.py',
        'plat': ['cuda']
    },
    'prospcd2sms': {
        'info': 'Converts ROS2 PointCloud2 topic to SMS memory_msgs::PointCloud Node',
        'app': 'spirecv/fromros2/ROS2PCD2SMSNode.py',
        'plat': ['cuda']
    },
    'paruco': {
        'info': 'Python Aruco QR-Code Detector',
        'app': 'spirecv/algorithm/aruco_det/ArucoDetNode.py',
        'plat': []
    },
    'caruco': {
        'info': 'C++ Aruco QR-Code Detector',
        'app': 'spirecv/algorithm/aruco_det/ArucoDetCppNode/build/SV2ArucoDetDemo',
        'plat': []
    },
    'cellipse': {
        'info': 'C++ Ellipse Detector',
        'app': 'spirecv/algorithm/ellipse_det/EllipseDetCppNode/build/SV2EllipseDetCppNode',
        'plat': []
    },
    'punderwaterenh': {
        'info': 'Nvidia-CUDA UnderWater Image Enhancement UDnet',
        'app': 'spirecv/algorithm/underwater_enhance/UDnet/UnderwaterUDnetNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyolo11rk': {
        'info': 'Python Rockchip Platform YOLOv11 Detector',
        'app': 'spirecv/algorithm/common_det/Rknn/YOLOv11DetNode_Rknn.py',
        'plat': ['rknn']
    },
    'cyolo11rk': {
        'info': 'C++ Rockchip Platform YOLOv11 Detector',
        'app': 'spirecv/algorithm/common_det/Rknn/YOLOv11DetCppNode_Rknn/build/rknn_yolo11_demo',
        'plat': ['rknn']
    },
    'pyolo11cuda': {
        'info': 'Nvidia-CUDA YOLOv11 Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11DetNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyoloworldcuda': {
        'info': 'Nvidia-CUDA YOLO-World Detector',
        'app': 'spirecv/algorithm/common_det/YOLOWorldNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyl11obbcuda': {
        'info': 'Nvidia-CUDA YOLOv11 Oriented Bounding Box (OBB) Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11ObbNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyl11posecuda': {
        'info': 'Nvidia-CUDA YOLOv11 Pose Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11PoseNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyl11segcuda': {
        'info': 'Nvidia-CUDA YOLOv11 Instance Segmentation Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11InstanceSegNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyolo8cuda': {
        'info': 'Nvidia-CUDA YOLOv8 Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv8DetNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyolo5cuda': {
        'info': 'Nvidia-CUDA YOLOv5-v7 Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv5v7DetNode_Cuda.py',
        'plat': ['cuda']
    },
    'pmmdetcuda': {
        'info': 'Nvidia-CUDA MMDetection Detector',
        'app': 'spirecv/algorithm/common_det/MMDet3Node_Cuda.py',
        'plat': ['cuda']
    },
    'pptpillarscuda': {
        'info': 'Nvidia-CUDA Point Pillars 3D Object Detector (Lidar-based)',
        'app': 'spirecv/algorithm/3d_det/PointPillars3DDetNode_Cuda.py',
        'plat': ['cuda']
    },
    'prtdetrcuda': {
        'info': 'Nvidia-CUDA RT-DETR Detector',
        'app': 'spirecv/algorithm/common_det/RTDETRDetNode_Cuda.py',
        'plat': ['cuda']
    },
    'pyolo11trt': {
        'info': 'Nvidia-TensorRT YOLOv11 Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11DetNode_Trt.py',
        'plat': ['trt']
    },
    'pyolo11segtrt': {
        'info': 'Nvidia-TensorRT YOLOv11 Instance Segmentation Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11InstanceSegNode_Trt.py',
        'plat': ['trt']
    },
    'pyolo11mbtrt': {
        'info': 'Nvidia-TensorRT YOLOv11 Multi-Batch Detector',
        'app': 'spirecv/algorithm/common_det/YOLOv11MultiBatchDetNode_Trt.py',
        'plat': ['trt']
    },
    'cbytetrackmot': {
        'info': 'ByteTrack Multiple Object Tracking Node',
        'app': 'spirecv/algorithm/mot/MultipleObjectTrackerCppNode/build/SV2MOTDemoMem',
        'plat': []
    },
    'pnanotrackrk': {
        'info': 'NanoTrack Single Object Tracking Node (Rockchip)',
        'app': 'spirecv/algorithm/sot/NanoTrack_Rknn/NanoTrackNode_Rknn.py',
        'plat': ['rknn']
    },
    'pocvsot': {
        'info': 'OpenCV-based Single Object Tracking Node',
        'app': 'spirecv/algorithm/sot/SingleObjectTrackerNode.py',
        'plat': []
    },
    'pclicktrackctl': {
        'info': 'Object Click Tracking Controller Node',
        'app': 'spirecv/algorithm/sot/ClickTrackControlNode.py',
        'plat': []
    },
    'cautofocus': {
        'info': 'Auto Focus Detector Node',
        'app': 'spirecv/algorithm/autofocus_det/AutoFocusDetCppNode/build/SV2AutoFocusDemo',
        'plat': []
    },
    'pmbsamcuda': {
        'info': 'Mobile Segment Anything Model Node (Nvidia-CUDA)',
        'app': 'spirecv/algorithm/mobilesam/MobileSAMNode_Cuda.py',
        'plat': ['cuda']
    },
    'pmbsamrk': {
        'info': 'Mobile Segment Anything Model Node (Rockchip)',
        'app': 'spirecv/algorithm/mobilesam/MobileSAMNode_Rknn.py',
        'plat': ['rknn']
    },
    'pcliprk': {
        'info': 'Clip Vision-Language Model Node (Rockchip)',
        'app': 'spirecv/algorithm/clip/ClipNode_Rknn.py',
        'plat': ['rknn']
    },
    'pcocoeval': {
        'info': 'COCO Dataset Evaluation Node',
        'app': 'spirecv/evaluator/COCOEvalNode.py',
        'plat': []
    },
    'pcocoassb': {
        'info': 'COCO Evaluation Data Aggregation Node',
        'app': 'spirecv/evaluator/COCOResultAssembleNode.py',
        'plat': []
    },
    'prtmpstm': {
        'info': 'RTMP Streaming Node (FFMpeg)',
        'app': 'spirecv/streamer/RTMPStreamerNode.py',
        'plat': []
    },
    'prtspstm': {
        'info': 'RTSP Streaming Node (FFMpeg)',
        'app': 'spirecv/streamer/RTSPStreamerNode.py',
        'plat': ['cuda']
    },
    'pgstmcuda': {
        'info': 'GStreaming Node (Gstreamer)',
        'app': 'spirecv/streamer/GStreamerNode_Cuda.py',
        'plat': ['cuda']
    },
    'crtspstmrk': {
        'info': 'RTSP Streaming Node (Rockchip)',
        'app': 'spirecv/streamer/StreamerCppNode_Rknn2/build/rtspstream',
        'plat': ['rknn']
    },
    'ertspserver': {
        'info': 'RTSP Streaming Server (ARM64)',
        'app': 'spirecv/streamer/StreamerCppNode_Rknn2/rtsp-simple-server_linux_arm64v8/rtsp-simple-server',
        'plat': ['arm']
    },
    'estmserveramd': {
        'info': 'RTSP/RTMP Streaming Server (X86_64)',
        'app': 'spirecv/streamer/ServerX86_64/mediamtx',
        'plat': ['x86']
    },
    'psvsave': {
        'info': 'Spire Format Image Result Saving Node',
        'app': 'spirecv/saver/SpireViewSaverNode.py',
        'plat': []
    },
    'csvjsave': {
        'info': 'Spire Format Video & SvjFile Saving Node',
        'app': 'spirecv/saver/SaverCppNode/build/SV2SaverDemo',
        'plat': []
    },
    'csvjsaverk': {
        'info': 'Spire Format Video & SvjFile Saving Node (Rockchip)',
        'app': 'spirecv/saver/SaverCppNode_Rknn2/build/SV2SaverDemo',
        'plat': ['rknn']
    },
    'pdetvis': {
        'info': 'Perception Result Visualization Node',
        'app': 'spirecv/visualizer/DetectionVisNode.py',
        'plat': []
    },
    'pnuvis': {
        'info': 'NuScenes Dataset & Result Visualization Node',
        'app': 'spirecv/visualizer/NuVisNode.py',
        'plat': []
    },
    'pnu2foxglove': {
        'info': 'NuScenes Dataset & Result To Foxglove Node',
        'app': 'spirecv/visualizer/NuFoxgloveVisNode.py',
        'plat': []
    },
    'ptgt2ros2': {
        'info': 'spirecv_msgs::2DTargets to ROS2 Targets Converting Node',
        'app': 'spirecv/toros2/SMS2DTargetsToROS2.py',
        'plat': []
    },
    'cgimbal': {
        'info': 'Gimbal Driver Node',
        'app': 'spirecv/gimbal/GimbalCppNode/build/SV2GimbalDemo',
        'plat': []
    },
    'cgbkeyctrl': {
        'info': 'Gimbal Keyboard Control Node',
        'app': 'spirecv/gimbal/GimbalKeyboardControlCppNode/build/SV2GimbalKeyboardControlDemo',
        'plat': []
    },
    'cgbtrack': {
        'info': 'Gimbal Tracking Node',
        'app': 'spirecv/gimbal/GimbalTrackingCppNode/build/SV2GimbalTrackingDemo',
        'plat': []
    },
    # Internal Testing
    'pyolodataload': {
        'info': 'Data Loading Pipeline Node in YOLOv11 Training',
        'app': 'spirecv/training/YOLO11TrainLoaderNode.py',
        'plat': ['test']
    },
    'pobjsizestat': {
        'info': 'Object Size Statistics Node',
        'app': 'spirecv/visualizer/ObjectSizeStatNode.py',
        'plat': ['test']
    },
    'ponnx2rk': {
        'info': 'YOLO11 Onnx To Rknn (Rockchip)',
        'app': 'spirecv/algorithm/common_det/Rknn/ModelOnnx2RknnNode.py',
        'plat': ['test']
    },
    'pwts2engine': {
        'info': 'YOLO11 Wts To Engine (TensorRT)',
        'app': 'spirecv/algorithm/common_det/YOLOv11DetJetsonTensorRTLib/Wts2EngineNode_Trt.py',
        'plat': ['test']
    }
}


def check_cuda_support() -> Dict[str, bool]:
    """检查是否支持 CUDA (NVIDIA GPU)"""
    cuda_support = False
    cuda_version = None
    try:
        # 方法1: 通过 nvidia-smi 命令检测
        result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=False)
        if "NVIDIA-SMI" in result.stdout:
            cuda_support = True
        # 方法2: 通过 PyTorch 检测（需安装 PyTorch）
        """
        try:
            import torch
            if torch.cuda.is_available():
                cuda_support = True
                cuda_version = torch.version.cuda
        except ImportError:
            pass
        """
    except Exception:
        pass
    return {"supported": cuda_support, "version": cuda_version}


def check_tensorrt_support() -> Dict[str, bool]:
    """检查是否支持 TensorRT"""
    tensorrt_support = False
    trt_version = None
    try:
        import tensorrt as trt
        tensorrt_support = True
        trt_version = trt.__version__
    except ImportError:
        pass
    return {"supported": tensorrt_support, "version": trt_version}


def check_openvino_support() -> Dict[str, bool]:
    """检查是否支持 Intel OpenVINO"""
    openvino_support = False
    ov_version = None
    try:
        from openvino.runtime import Core
        openvino_support = True
        ov_version = Core().get_versions("CPU")[0].description
    except ImportError:
        pass
    return {"supported": openvino_support, "version": ov_version}


def check_rknn_support() -> Dict[str, bool]:
    """检查是否支持瑞芯微 RKNN"""
    rknn_support = False
    rknn_version = None
    try:
        from rknn.api import RKNN
        rknn_support = True
        rknn_version = RKNN().version  # 实际需连接硬件设备
    except (ImportError, RuntimeError):
        pass
    return {"supported": rknn_support, "version": rknn_version}


def check_cann_support() -> Dict[str, bool]:
    """检查是否支持华为 CANN（昇腾 AI）"""
    cann_support = False
    cann_version = None
    try:
        import acl
        cann_support = True
        # 华为 CANN 版本通常通过环境变量获取
        cann_version = os.environ.get("ASCEND_VERSION", "unknown")
    except ImportError:
        pass
    return {"supported": cann_support, "version": cann_version}


def check_architecture():
    machine = platform.machine().lower()

    x86_keys = ('x86', 'i386', 'i686', 'amd64')
    if any(key in machine for key in x86_keys):
        return "x86"

    arm_keys = ('arm', 'aarch64', 'arm64')
    if any(key in machine for key in arm_keys):
        return "ARM"

    return f"Unknown Architecture ({machine})"


def check_os():
    system = platform.system()
    if system == "Linux":
        return "Linux"
    elif system == "Windows":
        return "Windows"
    elif system == "Darwin":
        return "macOS"
    else:
        return f"Other OS ({system})"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'program',
        type=str,
        default='list',
        help='RUN your SpireCV program. (INPUT "ls" to LIST available programs)')
    args, unknown_args = parser.parse_known_args()
    current_path = os.path.abspath(__file__)
    launch_fn = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'launch')
    launch_lst = []
    for f in os.listdir(launch_fn):
        if f.endswith('.launch'):
            launch_lst.append(Path(f).stem)
    launch_lst = sorted(launch_lst)

    if args.program in node_dict:
        program_fn = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], node_dict[args.program]['app'])
        """
        if len(unknown_args) > 0:
            program_fn += ' ' + ' '.join(unknown_args)
        """
        if program_fn.endswith('.py'):
            subprocess.run(["python", program_fn, *unknown_args])
        elif program_fn.endswith('rtsp-simple-server'):
            print(program_fn[:-len('rtsp-simple-server')] + 'rtsp-simple-server.yml')
            subprocess.run([program_fn, program_fn[:-len('rtsp-simple-server')] + 'rtsp-simple-server.yml'])
        elif program_fn.endswith('mediamtx'):
            print(program_fn[:-len('mediamtx')] + 'mediamtx.yml')
            subprocess.run([program_fn, program_fn[:-len('mediamtx')] + 'mediamtx.yml'])
        else:
            subprocess.run([program_fn, *unknown_args])
    elif args.program in launch_lst or args.program.endswith('.launch'):
        if args.program.endswith('.launch'):
            launch_f = args.program
        else:
            launch_f = os.path.join(launch_fn, args.program + '.launch')
        with open(launch_f, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        procs = []
        for line in lines:
            line = line.strip()
            if line.startswith('#'):
                continue
            if line.startswith('smsrun '):
                line = line[7:]
            if len(line) > 0:
                cmd = line.split(' ')
                if 'sleep' == cmd[0]:
                    sleep_sec = float(cmd[1])
                    print("sleep in {} sec".format(sleep_sec))
                    time.sleep(sleep_sec)
                    continue
                assert cmd[0] in node_dict
                cmd[0] = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], node_dict[cmd[0]]['app'])
                if cmd[0].endswith('.py'):
                    cmd.insert(0, 'python')
                elif cmd[0].endswith('rtsp-simple-server'):
                    cmd.append(cmd[0][:-len('rtsp-simple-server')] + 'rtsp-simple-server.yml')
                elif cmd[0].endswith('mediamtx'):
                    cmd.append(cmd[0][:-len('mediamtx')] + 'mediamtx.yml')
                print(cmd)
                proc = subprocess.Popen([*cmd])
                procs.append(proc)
        try:
            for p in procs:
                p.wait()
        except KeyboardInterrupt:
            for p in procs:
                if p.poll() is None:
                    p.kill()
                    p.wait()
    elif args.program in ['list', 'ls']:
        key_word = None
        list_inner = False
        if len(unknown_args) > 0:
            if unknown_args[0] == '-a':
                list_inner = True
            else:
                key_word = unknown_args[0]
        width = 0
        print("\033[32;1mApp:\033[0m")
        for key in node_dict:
            width = max(width, len(key))
        max_w = max(20, shutil.get_terminal_size().columns)
        for key, val in node_dict.items():
            if key_word is not None and key_word not in key:
                continue
            if not list_inner and 'test' in val['plat']:
                continue
            str_show_orin = "  {} {} -- [{}]".format(key.ljust(width) + ":", val['info'], val['app'])
            str_show = "  \033[1m{}\033[0m: {} -- [{}]".format(key.ljust(width), val['info'], val['app'])
            if len(str_show_orin) > max_w:
                str_show = str_show[:max_w + (len(str_show) - len(str_show_orin)) - 2] + '..'
            print(str_show)
        print("\033[32;1mLaunch:\033[0m")
        for lch in launch_lst:
            if key_word is not None and key_word not in lch:
                continue
            print("  \033[1m{}\033[0m".format(lch))
    else:
        print('The program does not exist.')


if __name__ == '__main__':
    main()
